#include "pktfwd.h"
#include "obuf.h"
#include "ibuf.h"

using namespace ai_sim;

Define_Module(Pktfwd);

void Pktfwd::initialize() 
{
	// normal parameters
	numPorts = par("numPorts");
	is_cross_connected = par("is_cross_connected");
	cur_switch_index = getParentModule()->getIndex(); 
	is_HD_mode = par("is_HD_mode");

	// only to be used in cross-connected network switch
	total_level_switch = par("total_level_switch");
	cross_connected_upperlevel_numswitch = par("cross_connected_upperlevel_numswitch");
	
	Switch = getParentModule();
	if (!Switch) 
	{
		error("-E- Failed to obtain an parent Switch module");
	}
}

void Pktfwd::handleMessage(omnetpp::cMessage *p_msg)
{   
	if (p_msg->getKind() == ROUTE_ALLOC_MSG)
	{
		int base_gate_id = getParentModule()->gate("port$o",0)->getId();
		RouteVecMsg *rv_msg = dynamic_cast<RouteVecMsg*> (p_msg);
		if (rv_msg == NULL)
		{
			error("Fail in dynamic casting incoming msg to Route Vector message in pktfwd!!");
		}
		routingtable.clear();
		std::map<int,int>::iterator It;
		for (auto &elm:rv_msg->getNew_route_table())
		{
			It = routingtable.find(elm.destIdx);
			if (It != routingtable.end()) // elm.first is the sLID
				// for the same LID there are multiple output ports option!!
				error("Fail in configuring routing table as encounter multiple output ports for LID:%d !!",elm.destIdx);
			else
				routingtable[elm.destIdx] = elm.outPort-base_gate_id; //store the output port num for given LID in routingtable
		}
		char path[40];
		sprintf(path,"sw_l%d",total_level_switch);
		static bool send_routing_to_upper_switch = true; // we only need to forward the routing info from the first switch to other switches once		
		if (!strcmp(getParentModule()->getName(),path) && send_routing_to_upper_switch ) 
		// for highest level switch in cross-connected level, the first switch need to broadcast the routing 
		// info it received from central controller to the other switches
		{
		      	for(int i = 1; i < int(numPorts/2); i++)
		      	{
				char path[40];
				sprintf(path,"^.^.sw_l%d[%d].pktfwd",total_level_switch,i);
				Pktfwd* temp_pktfwd = dynamic_cast<Pktfwd*>(getModuleByPath(path));
				RouteVecMsg *new_rv_msg = rv_msg->dup();
				sendDirect(new_rv_msg,temp_pktfwd,"routing_info_port");
		      	}
		      	send_routing_to_upper_switch = false;
		}
		
		for (auto temp:routingtable)
		{
			EV<<temp.first<<" "<<temp.second<<omnetpp::endl;
		}
		delete rv_msg;
	}
	else
		error(" Received message type: %d non-ROUTE_ALLOC_MSG type on pktfwd!!",p_msg->getKind());
}

// get the output port for the given LID - the actual AR or deterministic routing
int Pktfwd::getOutputPort(IBDataMsg *p_dataMsg) 
{
	unsigned int dLid = p_dataMsg->getDstLid();
	unsigned int virtual_lane = p_dataMsg->getVL();
	unsigned int outPort; // the resulting output port
	Enter_Method("getPortByLID LID: %d", dLid);
	std::map<int,int>::iterator It;
	It = routingtable.find(dLid);
	if (!is_cross_connected) // current switch is not cross-connected to its upper level switch
	{
		char path[40];
		sprintf(path,"sw_l%d",total_level_switch);
		if (strcmp(getParentModule()->getName(),"core") && strcmp(getParentModule()->getName(),path) || is_HD_mode )
		// in the case of non-highest level switch or
		{
			if (It != routingtable.end())
			{	
				EV<<"DLID:"<<dLid<<"portout:"<<routingtable.at(dLid)<<omnetpp::endl;
				return routingtable.at(dLid);
			}
			else// if we cant get DLID from routingtable, meaning the dest GPU is not connected to current switch,
				// must forward it to upper level switch through the last output port
			{
				EV<<"DLID:"<<dLid<<"portout:"<<numPorts<<omnetpp::endl;
				return numPorts; 
			}
		}
		else // in the case of highest level switch
		{
			if (It != routingtable.end())
			{	
				EV<<"DLID:"<<dLid<<"portout:"<<routingtable.at(dLid)<<omnetpp::endl;
				return routingtable.at(dLid);
			}
			int previous = routingtable.begin()->second;	
			for (auto temp:routingtable)
			{
				if (dLid < temp.first) // check which category the DLID belongs to
				{	
					std::cout<<"DLID:"<<dLid<<"portout:"<<previous<<omnetpp::endl;
					return previous;
				}	
				previous = temp.second;
			}
			EV<<"DLID:"<<dLid<<"portout:"<<routingtable.rbegin()->second<<omnetpp::endl; 
			return routingtable.rbegin()->second;
		}
	}
	else // current switch is cross-connected to it's upper level switch
	{
		if (It != routingtable.end())
		{	
			EV<<"DLID:"<<dLid<<"portout:"<<routingtable.at(dLid)<<omnetpp::endl;
			return routingtable.at(dLid);
		}
		else// if we cant get DLID from routingtable, then we have to forward it to upper level switch
		{
			int most_free_credit_available = -1;
			int most_free_credit_available_switch = -1;
			int base_index = getParentModule()->gateSize("port$o") - cross_connected_upperlevel_numswitch;
			IBInBuf* temp_ibuf;
			for (int i = 0; i<=int(cross_connected_upperlevel_numswitch/2);i++) 
			// we begin from upper level switch with index:cur_switch_index/2, followed by the adjacent switches
			// the threshold is set at (maxcredit /8), if the available credit on target switch is less than maxcredit/8, we move on to adjacent switch
			// worst case is every upper level switch's available credit is below the threshold, then we pick the switch with the best traffic
			{
				// left neighbour switch
				int upperlevel_switch_index_left = (int(cur_switch_index/2)+cross_connected_upperlevel_numswitch+i)%cross_connected_upperlevel_numswitch;
				char path[40];
				sprintf(path,"^.^.sw_l%d[%d].subport[%d].ibuf",total_level_switch,upperlevel_switch_index_left,cur_switch_index);
				temp_ibuf = dynamic_cast<IBInBuf*>(getModuleByPath(path));
				if (temp_ibuf->staticFree.at(virtual_lane) > int(temp_ibuf->maxStatic.at(virtual_lane)/8))
				{
	EV<<"DLID:"<<dLid<<"portout:"<<base_index+upperlevel_switch_index_left<<omnetpp::endl;
					return base_index+upperlevel_switch_index_left;
				}
				
				if (temp_ibuf->staticFree.at(virtual_lane) > most_free_credit_available)
				{
					most_free_credit_available = temp_ibuf->staticFree.at(virtual_lane);
					most_free_credit_available_switch = upperlevel_switch_index_left;
				}
				
				
				// right neighbour switch
				int upperlevel_switch_index_right = (int(cur_switch_index/2)+cross_connected_upperlevel_numswitch-i)%cross_connected_upperlevel_numswitch;
				sprintf(path,"^.^.sw_l%d[%d].subport[%d].ibuf",total_level_switch,upperlevel_switch_index_left,cur_switch_index);
				temp_ibuf = dynamic_cast<IBInBuf*>(getModuleByPath(path));
				if (temp_ibuf->staticFree.at(virtual_lane) > int(temp_ibuf->maxStatic.at(virtual_lane)/8))
				{
					EV<<"DLID:"<<dLid<<"portout:"<<base_index+upperlevel_switch_index_right<<omnetpp::endl;
					return base_index+upperlevel_switch_index_right;
				}
				
				if (temp_ibuf->staticFree.at(virtual_lane) > most_free_credit_available)
				{
					most_free_credit_available = temp_ibuf->staticFree.at(virtual_lane);
					most_free_credit_available_switch = upperlevel_switch_index_right;
				}
			}
			// we return the outport number corresponding to most suitable (best traffic) swicth 
			EV<<"DLID:"<<dLid<<"portout:"<<base_index+most_free_credit_available_switch<<omnetpp::endl;
			return base_index+most_free_credit_available_switch;
		}
	}
}

// report queuing of flits on TQ for DLID (can be negative for arb)
int Pktfwd::repQueuedFlits(unsigned int rq, unsigned int tq, unsigned int dlid, int numFlits) {
	Enter_Method("repQueuedFlits tq:%d flits:%d", tq, numFlits);
	return(0);
}

// IBuf received a TQLoadUpdate - Handle Received Port Usage Notification
void Pktfwd::handleTQLoadMsg(unsigned int tq, unsigned int srcRank, unsigned int firstLid, unsigned int lastLid, int load) {
	Enter_Method("handleTQLoadMsg tq:%d srcRank:%d lid-range: [%d,%d] load:%d", tq, srcRank, firstLid,
			lastLid, load);
	EV << "-I- " << getFullPath() << " handleTQLoadMsg tq: " << tq << " srcRank: " << srcRank << " lids: "
			<< firstLid << "," << lastLid << " load: " << load << omnetpp::endl;

}

void Pktfwd::finish()
{
}

Pktfwd::~Pktfwd() 
{
}